package com.atguigu.bigdata.spark.core.rdd.partition

import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkConf, SparkContext}

/**
  * create by undeRdoG on  2021-06-14  17:28
  * 凡心所向，素履以往，生如逆旅，一苇以航。
  */
object Spark01_RDD_Part {

  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setAppName("Partition").setMaster("local[*]")

    val sc = new SparkContext(sparkConf)

    val rdd: RDD[(String, String)] = sc.makeRDD(List(
      ("nba", "xxxxxxx"),
      ("cba", "xxxxxxxxxxxx"),
      ("wnba", "xxxxxxxxxxxxx")
    ))
    val res: RDD[(String, String)] = rdd.partitionBy(new MyPartitioner)

    res.saveAsTextFile("partition/output")

  }

  class MyPartitioner extends Partitioner{
    override def numPartitions: Int = 3

    // 返回数据的分区索引  从0 开始
    override def getPartition(key: Any): Int = {
      if (key == "nba"){
        0
      }else if (key == "wnba"){
        1
      }else{
        2
      }
    }
  }

}
